#!/home/yangzheng/anaconda3/bin/python
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import pyplot as plt
from typing import Callable
from scipy.stats import norm

TIME_GAP = 0.1


def my_softmax(x):
    exp_x = np.exp(x)
    return exp_x / np.sum(exp_x)


class ParticleFilter:
    def __init__(self, particles, propagate: Callable, origin_observation):
        self.propagate = propagate
        self.particles = particles
        self.belief = np.ones(len(particles)) / (len(particles))
        self.previous_observation = origin_observation
        self.posterior_probability_of_current_observation = 0

    # this function update the belief of every particle
    def updateBelief(self, observation):
        self.posterior_probability_of_current_observation = 0
        # calculate the posterior probability density under the condition of particle[i]
        for i, particle_belief in enumerate(self.belief):
            predicted_distribution, predicted_position = self.propagate(self.previous_observation, self.particles[i])
            # calculate distance between the predicted point and ground truth
            distance = ((predicted_position[0] - observation[0]) ** 2 + (predicted_position[1] - observation[1]) ** 2) ** 0.5
            # update the particle belief based on bayesian equation
            particle_belief = particle_belief * (1 / (distance**2 + 0.0001))
            # particle_belief = particle_belief * (post_prob_density_of_current_observation * 100)
            self.belief[i] = particle_belief


        # normalize belief
        # add an offset to prevent particle loss
        offset = 0
        self.belief = self.belief + offset
        sum = np.sum(self.belief)
        self.belief = self.belief / sum

        # update observation
        self.previous_observation = observation


# This part is to be substituted by a driver model
def propagate(x, particle: float):
    time_gap = TIME_GAP
    scale = 0.05

    def distributionFunction(x_new):
        # here particle means speed
        return norm.pdf(x_new, x + particle * time_gap, scale)

    return distributionFunction


if __name__ == '__main__':
    particles = np.linspace(5, 50, 200)
    origin_position = 0.0
    his_filter = ParticleFilter(particles=particles, propagate=propagate, origin_observation=origin_position)
    position_list = np.array([1])
    delta_x = 0.1
    plt.ion()
    for i in range(0, 80):
        position_list = np.append(position_list, position_list[i] + 1 + delta_x)
        if i <= 30:
            delta_x += 0.05
        if i >= 60:
            delta_x -= 0.05
    print(position_list[-1] - position_list[-2])
    for i, observe in enumerate(position_list):
        his_filter.updateBelief(observe)
        # plot the current result
        speed_inferred = particles[np.argmax(his_filter.belief)]
        actual_speed = (position_list[i + 1] - position_list[i]) / TIME_GAP
        # plt.ylim(0.01, 0.08)
        plt.plot(particles, his_filter.belief)
        plt.text(30, plt.ylim()[1] * 0.8, 'inferred speed: ' + str(speed_inferred))
        plt.text(30, plt.ylim()[1] * 0.7, 'actual speed: ' + str(actual_speed))
        plt.pause(0.1)
        plt.clf()
    plt.show()
